from apscheduler.schedulers.blocking import BlockingScheduler
from ssq_crawler import SSQCrawler
from dlt_crawler import DLTCrawler
from database import Database
import logging
from datetime import datetime
import os

class LotteryDataCollector:
    def __init__(self):
        self.ssq_crawler = SSQCrawler()
        self.dlt_crawler = DLTCrawler()
        self.db = Database()
        self.scheduler = BlockingScheduler()

    def collect_data(self):
        """收集并保存开奖数据"""
        try:
            # 爬取双色球数据
            ssq_data = self.ssq_crawler.crawl()
            if ssq_data:
                ssq_count = self.db.save_draws(ssq_data, 'ssq')
                logging.info(f'Successfully saved {ssq_count} SSQ lottery records')
            else:
                logging.warning('No SSQ data collected')

            # 爬取大乐透数据
            dlt_data = self.dlt_crawler.crawl()
            if dlt_data:
                dlt_count = self.db.save_draws(dlt_data, 'dlt')
                logging.info(f'Successfully saved {dlt_count} DLT lottery records')
            else:
                logging.warning('No DLT data collected')

        except Exception as e:
            logging.error(f'Error in data collection: {str(e)}')

    def start_scheduler(self):
        """启动定时任务"""
        # 每天22:00执行一次数据收集
        self.scheduler.add_job(
            self.collect_data,
            'cron',
            hour=22,
            minute=0,
            id='collect_lottery_data'
        )

        try:
            logging.info('Starting scheduler...')
            self.scheduler.start()
        except (KeyboardInterrupt, SystemExit):
            self.stop()

    def stop(self):
        """停止定时任务并清理资源"""
        self.scheduler.shutdown()
        self.db.close()
        logging.info('Collector stopped')

def main():
    # 设置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.StreamHandler(),
            logging.FileHandler('lottery_collector.log')
        ]
    )

    collector = LotteryDataCollector()
    
    # 首次运行立即收集一次数据
    collector.collect_data()
    
    # 启动定时任务
    # collector.start_scheduler()

if __name__ == '__main__':
    main()